from liblinearutil import *
import random


def calcu(true_tag, pre_tag):
    true_num = {}
    pre_num = {}
    corr_num = {}
    for iT, iP in zip(true_tag, pre_tag):
        if iT not in true_num:
            true_num[iT] = 0
            pre_num[iT] = 0
            corr_num[iT] = 0
        if iP not in true_num:
            true_num[iP] = 0
            pre_num[iP] = 0
            corr_num[iP] = 0
        true_num[iT] += 1
        pre_num[iP] += 1

        if iP == iT:
            corr_num[iP] += 1
    precision = []
    recall = []
    fmeasure = []

    for key in true_num:
        weight = true_num[key] / len(true_tag)
        tm1 = corr_num[key] / pre_num[key]
        tm2 = corr_num[key] / true_num[key]
        precision.append(tm1 * weight)
        recall.append(tm2 * weight)
        fmeasure.append(2 * tm1 * tm2 / (tm1 + tm2) * weight)
    print(precision)
    print(recall)
    print(fmeasure)
    return sum(precision), sum(recall), sum(fmeasure)


def random_data(data, tag, testIds):
    train_data = []
    train_tag = []
    test_data = []
    test_tag = []
    ptests = set(testIds)

    for i in range(len(data)):
        if i in ptests:
            test_data.append(data[i])
            test_tag.append(tag[i])
        else:
            train_data.append(data[i])
            train_tag.append(tag[i])

    return train_data, train_tag, test_data, test_tag


def random_id(y, num):
    l = [i for i in range(y)]
    random.shuffle(l)
    con = int(y / num)
    pl = []
    for i in range(1, num):
        pl.append(sorted(l[(i - 1) * con:i * con]))
    pl.append(sorted(l[con * (num - 1):]))

    return pl


def predicts(y, x, n_fold):
    pre_tag = [-1 for i in range(len(y))]
    alltest = random_id(len(y), n_fold)

    for testids in alltest:
        train_data, train_tag, test_data, test_tag = random_data(x, y, testids)
        m = train(train_tag, train_data)
        p_label, p_acc, p_val = predict(test_tag, test_data, m)
        for pindex, p_one_label in enumerate(p_label):
            pre_tag[testids[pindex]] = p_one_label

    print(calcu(y, pre_tag))


if __name__ == "__main__":
    num = 5
    y, x = svm_read_problem('../data/new-tweet.data.idf_svm')
    # m = train(y, x,"-v 5")
    predicts(y, x, 5)
